from starlette.applications import Starlette
from starlette.routing import Mount
from mcp.server.fastmcp import FastMCP
import contextlib

echo_mcp = FastMCP(name="EchoServer", stateless_http=True)


@echo_mcp.tool()
def echo(message: str) -> str:
    return f"Echo:{message}"


math_mcp = FastMCP(name="MathServer", stateless_http=True)


@math_mcp.tool()
def add_two(n: int) -> int:
    return n + 2


# 定义Starlette的应用的生命周期管理器，负责启动和关闭两个MCP会话管理器
@contextlib.asynccontextmanager
async def lifespan(app: Starlette):
    # 创建异步退出栈，确保多个上下文正常的进入和退出
    async with contextlib.AsyncExitStack() as stack:
        # 启动echo_mcp的会话管理
        await stack.enter_async_context(echo_mcp.session_manager.run())
        # 启动math_mcp的会话管理
        await stack.enter_async_context(math_mcp.session_manager.run())
        # 生命周期yield 供Starlette使用
        yield


app = Starlette(
    routes=[
        # 挂载echo_mcp到/echo路径，实际的接口为/echo/mcp
        Mount("/echo", app=echo_mcp.streamable_http_app()),
        # 挂载math_mcp到/math路径，实际的接口为/math/mcp
        Mount("/math", app=math_mcp.streamable_http_app()),
    ],
    # 指定生命周期管理器
    lifespan=lifespan,
)


# 定义主函数
def main():
    # 导入 uvicorn 服务器
    import uvicorn

    # 启动应用服务器
    uvicorn.run(app, host="127.0.0.1", port=8000)


# 脚本入口
if __name__ == "__main__":
    main()
# uvicorn asgi_server:app --host 127.0.0.1 --port 8000 --reload
# curl -X POST http://localhost:8000/echo/mcp \
#  -H "Content-Type: application/json" \
# -H "Accept: application/json, text/event-stream" \
#  -d '{
#    "jsonrpc": "2.0",
#    "method": "tools/call",
#    "params": {
#      "name": "echo",
#      "arguments": {"message": "Hello World"}
#    },
#    "id": 1
#  }'
#
#
# curl -X POST http://localhost:8000/math/mcp \
#  -H "Content-Type: application/json" \
# -H "Accept: application/json, text/event-stream" \
#  -d '{
#    "jsonrpc": "2.0",
#    "method": "tools/call",
#    "params": {
#      "name": "add_two",
#      "arguments": {"n": 5}
#    },
#    "id": 1
#  }'
